import os
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import precision_score, recall_score, f1_score, accuracy_score, confusion_matrix
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from joblib import dump
# 数据加载
df = pd.read_csv('Normalization_total_extend.csv')
X = df.drop(columns=['Label'])
y = df['Label']

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=50)

# 检查模型文件是否存在，若存在则加载模型
model_file = 'RandomForestModel/random_forest_model.pkl'
if os.path.exists(model_file):
    rf_model = pd.read_pickle(model_file)
else:
    # 创建随机森林模型
    rf_model = RandomForestClassifier(n_estimators=100, random_state=50)

    # 训练模型
    rf_model.fit(X_train, y_train)

    # 保存模型
    #rf_model.to_pickle(model_file)
    dump(rf_model, model_file)
# 训练集上的性能评估
y_train_pred = rf_model.predict(X_train)
train_precision = precision_score(y_train, y_train_pred, average='weighted')
train_recall = recall_score(y_train, y_train_pred, average='weighted')
train_f1 = f1_score(y_train, y_train_pred, average='weighted')
train_accuracy = accuracy_score(y_train, y_train_pred)
print(f"Training Precision: {train_precision:.4f}")
print(f"Training Recall: {train_recall:.4f}")
print(f"Training F1-Score: {train_f1:.4f}")
print(f"Training Accuracy: {train_accuracy:.4f}")

# 测试集上的性能评估
y_test_pred = rf_model.predict(X_test)
test_precision = precision_score(y_test, y_test_pred, average='weighted')
test_recall = recall_score(y_test, y_test_pred, average='weighted')
test_f1 = f1_score(y_test, y_test_pred, average='weighted')
test_accuracy = accuracy_score(y_test, y_test_pred)
print(f"Test Precision: {test_precision:.4f}")
print(f"Test Recall: {test_recall:.4f}")
print(f"Test F1-Score: {test_f1:.4f}")
print(f"Test Accuracy: {test_accuracy:.4f}")

# 绘制混淆矩阵
cm = confusion_matrix(y_test, y_test_pred)
plt.figure(figsize=(10, 7))
plt.imshow(cm, cmap='Blues')
plt.title('Confusion Matrix')
plt.xlabel('Predicted')
plt.ylabel('Actual')
plt.xticks(range(cm.shape[1]), ['Class {}'.format(i) for i in range(1, cm.shape[1] + 1)], rotation=45)
plt.yticks(range(cm.shape[0]), ['Class {}'.format(i) for i in range(1, cm.shape[0] + 1)])
plt.colorbar()
plt.grid(False)
plt.show()

# 保存测试集上的性能评估结果图（如果需要）
plt.savefig('RandomForestModel/random_forest_confusion_matrix.png')